/***********************************
 Author       : LIN
 Email        : l2441473911@163.com
 LastModified : 2020-12-28 23:4
 FileName     : main.c
 Description  : main program to carry out
                encoding/decoding
***********************************/
#include <stdio.h>
#include <string.h>
#include "bit-stack.h"
#include "priority-queue.h"
#include "utility.h"
#include "main.h"

int main(int argc, char **argv){
	HandleParameters(argc, argv);
	return 0;
}

// freq_list should has the size of 95, each element of which
// indicates the appearance frequency of one single character
// ranging from 0x20 to 0x7E in ASCII (all the printable cha-
// racters).
HuffmanForest InitializeHuffmanForest(int *freq_list){
	int i, j, n;
	HuffmanForest forest;
	HuffmanTreeLeave *leaf, **leave_list;
	
	// more comfortable to me to know 'how many' first
	for(n = 0, i = 0; i < TOTAL_SYMBOL_NUMBER; i++)
		if(freq_list[i]) n++;

	if(n > 0){
		leave_list = (HuffmanTreeLeave **)malloc(n * sizeof(void *));

		for(i = 0, j = 0; i < TOTAL_SYMBOL_NUMBER; i++){
			if(freq_list[i]){
				leaf = (HuffmanTreeLeave *)malloc(sizeof(HuffmanTreeLeave));
				leaf->weight = freq_list[i];
				leaf->symbol = i + ASCII_PRINTABLE_CHAR_START;
				leaf->type = 1;
				leave_list[j++] = leaf;
			}
		}

		forest = BuildHeap((void **)leave_list, n, n + 1);
		return forest;
	}
	return NULL;
}

HuffmanTree *ConstructHuffmanTree(HuffmanForest forest){
	if(forest == NULL) return NULL;

	int i, n = forest->heap_size;
	void *leftTree, *rightTree;
	HuffmanTreeInteriorNode *newTree;
	HuffmanTree *completeTree;

	for(i = 1; i < n; i++){
		leftTree = HeapPopMin(forest);
		rightTree = HeapPopMin(forest);

		newTree = (HuffmanTreeInteriorNode *)malloc(sizeof(HuffmanTreeInteriorNode));
		newTree->weight = DerefAsInt(leftTree) + DerefAsInt(rightTree);
		newTree->leftChild = leftTree;
		newTree->rightChild = rightTree;
		newTree->type = 0;

		HeapInsert(forest, newTree);
	}

	completeTree = HeapPopMin(forest);

	return completeTree;
}

void _EncodeThoughHuffmanTree_Rec(void *tree, BitStack *stack, EncodeTable table){
	if(HuffTreeIsLeaf(tree)){
		HuffmanTreeLeave *leaf = (HuffmanTreeLeave *)tree;

		BitCode *code = (BitCode *)malloc(sizeof(BitCode));

		code->code_length = stack->top + 1;

		int leastByteNumber = stack->top / BITS_PRE_BYTE + 1;
		code->code_array = (unsigned char *)malloc(leastByteNumber);

		// clear code_array
		for(leastByteNumber--; leastByteNumber >= 0; leastByteNumber--)
			code->code_array[leastByteNumber] = 0x00u;

		// copy the current stack content to code_array from bottom to top
		BitStackCopyStackBottomToTopTo(stack, code->code_array);

		table[leaf->symbol - ASCII_PRINTABLE_CHAR_START] = code;

		BitStackPop(stack, NULL);
	}else{
		BitStackPush(stack, 0x00u);
		_EncodeThoughHuffmanTree_Rec(((HuffmanTreeInteriorNode *)tree)->leftChild, stack, table);
		BitStackPush(stack, 0x01u);
		_EncodeThoughHuffmanTree_Rec(((HuffmanTreeInteriorNode *)tree)->rightChild, stack, table);
		BitStackPop(stack, NULL);
	}
}

EncodeTable EncodeThroughHuffmanTree(HuffmanTree *tree){
	if(tree == NULL) return NULL;
	EncodeTable table = (EncodeTable)malloc(TOTAL_SYMBOL_NUMBER * sizeof(void *));
	// assume that the length our codeword will not exceeds the number of bits
	// in an [int]
	BitStack *stack = CreateBitStack(MAX_BITCODE_LENGTH);

	// clear table
	int i;
	for(i = 0; i < TOTAL_SYMBOL_NUMBER; i++)
		table[i] = NULL;

	_EncodeThoughHuffmanTree_Rec(tree, stack, table);

	DisposeBitStack(stack);

	return table;
}

void PrintHuffTree(void *t){
	if(HuffTreeIsLeaf(t)){
		fprintf(stdout, "%d(%c) ", ((HuffmanTreeLeave *)t)->symbol, ((HuffmanTreeLeave *)t)->symbol);
	}else{
		fprintf(stdout, "0 ");
		PrintHuffTree(((HuffmanTreeInteriorNode *)t)->leftChild);
		PrintHuffTree(((HuffmanTreeInteriorNode *)t)->rightChild);
	}
}

void ShowEncodeTable(EncodeTable table){
	int i;
	for(i = 0; i < TOTAL_SYMBOL_NUMBER; i++){
		if(table[i] != NULL){
			fprintf(stdout, "%c(%d): ", ASCII_PRINTABLE_CHAR_START + i, ASCII_PRINTABLE_CHAR_START + i);
			_ShowBits(table[i]->code_array, table[i]->code_length);
		}
	}
}

void WriteDecodeHuffmanTreeToFile(void *complete_huff_tree, FILE *fp){
	char is_leave = HuffTreeIsLeaf(complete_huff_tree);
	if(is_leave){
		fwrite(&(((HuffmanTreeLeave *)complete_huff_tree)->symbol), sizeof(char), 1, fp);
	}else{
		fwrite(&is_leave, sizeof(char), 1, fp);
		WriteDecodeHuffmanTreeToFile(((HuffmanTreeInteriorNode *)complete_huff_tree)->leftChild, fp);
		WriteDecodeHuffmanTreeToFile(((HuffmanTreeInteriorNode *)complete_huff_tree)->rightChild, fp);
	}
}

void *ReadDecodeHuffmanTreeFromFile(FILE *fp){
	char c;
	if(fread(&c, sizeof(char), 1, fp) == 1){
		if(c){
			DecodeHuffTreeLeaf *leaf = (DecodeHuffTreeLeaf *)malloc(sizeof(DecodeHuffTreeLeaf));
			leaf->type = 1;
			leaf->symbol = c;
			return leaf;
		}else{
			DecodeHuffTreeInteriorNode *node = (DecodeHuffTreeInteriorNode *)malloc(sizeof(DecodeHuffTreeInteriorNode));
			node->type = 0;
			node->leftChild = ReadDecodeHuffmanTreeFromFile(fp);
			node->rightChild = ReadDecodeHuffmanTreeFromFile(fp);
			return node;
		}
	}
	return NULL;
}

void PrintDecodeHuffTree(void *t){
	if(DecodeHuffTreeIsLeaf(t)){
		fprintf(stdout, "%d(%c) ", ((DecodeHuffTreeLeaf *)t)->symbol, ((DecodeHuffTreeLeaf *)t)->symbol);
	}else{
		fprintf(stdout, "0 ");
		PrintDecodeHuffTree(((DecodeHuffTreeInteriorNode *)t)->leftChild);
		PrintDecodeHuffTree(((DecodeHuffTreeInteriorNode *)t)->rightChild);
	}
}

void DisposeHuffmanForest(HuffmanForest f){
	DisposeHeap(f);
}

void DisposeHuffmanTree(void *t){
	if(!HuffTreeIsLeaf(t)){
		// any non-leaf subtree of a huffman tree
		// has two child
		DisposeHuffmanTree(((HuffmanTreeInteriorNode *)t)->leftChild);
		DisposeHuffmanTree(((HuffmanTreeInteriorNode *)t)->rightChild);
	}
	free(t);
}

void DisposeEncodeTable(EncodeTable table){
	int i;
	for(i = 0; i < TOTAL_SYMBOL_NUMBER; i++){
		if(table[i] != NULL){
			free(table[i]->code_array);
			free(table[i]);
		}
	}
	free(table);
}

void DisposeDecodeHuffTree(void *t){
	if(!DecodeHuffTreeIsLeaf(t)){
		DisposeDecodeHuffTree(((DecodeHuffTreeInteriorNode *)t)->leftChild);
		DisposeDecodeHuffTree(((DecodeHuffTreeInteriorNode *)t)->rightChild);
	}
	free(t);
}

void EncodeFile(FILE *fp_in, FILE *fp_out, EncodeTable encode_table){
	int byte_index, one_char;
	unsigned char one_byte, set_one_mask, check_bit_mask, bits_has_set, bits_has_used, end_of_input_file = 0;
	unsigned int all_bits_written = 0;
	BitCode *current_bitcode;
	
	fseek(fp_in, 0, SEEK_SET);
	fseek(fp_out, 0, SEEK_SET);

	one_char = fgetc(fp_in);
	current_bitcode = encode_table[one_char - ASCII_PRINTABLE_CHAR_START];
	if(one_char < ASCII_PRINTABLE_CHAR_START || one_char > ASCII_PRINTABLE_CHAR_END || current_bitcode == NULL){
		fprintf(stderr, "The file may contain symbol(s) that have no corresponding codeword.\n");
		return;
	}

	byte_index = 0;
	bits_has_used = 0;
	check_bit_mask = 0x80u;

	while(!end_of_input_file){
		// sets all 8 bits of one_byte and writes one_byte
		// to the output file in one iteration in the outer
		// loop

		one_byte = 0x00u;
		set_one_mask = 0x80u;
		bits_has_set = 0;
		
		// sets all 8 bits of one_byte
		while(bits_has_set < BITS_PRE_BYTE){
			// set one only, since we have initialize
			// one_byte to all zero.
			if(current_bitcode->code_array[byte_index] & check_bit_mask)
				one_byte |= set_one_mask;

			// updates variables accordingly
			bits_has_set++;
			all_bits_written++;
			set_one_mask >>= 1;
			
			if(bits_has_used + 1 == current_bitcode->code_length){
				one_char = fgetc(fp_in);
				// skip return(0x0D) and newline(0x0A)
				// since we do not encode these symbols
				// but they will appear in the end of a
				// text file even if these is only one
				// line of text in the file.
				while(one_char == 0x0A || one_char == 0x0D)
					one_char = fgetc(fp_in);
				
				if(one_char == EOF){
					end_of_input_file = 1;
					break;
				}
				current_bitcode = encode_table[one_char - ASCII_PRINTABLE_CHAR_START];
				if(one_char < ASCII_PRINTABLE_CHAR_START || one_char > ASCII_PRINTABLE_CHAR_END || current_bitcode == NULL){
					fprintf(stderr, "The file may contain symbol(s) that have no corresponding codeword.\n");
					return;
				}
				byte_index = 0;
				bits_has_used = 0;
				check_bit_mask = 0x80u;
			}else if(check_bit_mask == 0x01u){
				byte_index++;
				bits_has_used++;
				check_bit_mask = 0x80u;
			}else{
				bits_has_used++;
				check_bit_mask >>= 1;
			}
		}

		fwrite(&one_byte, sizeof(unsigned char), 1, fp_out);
	}

	// write the information of how many bits should be examined
	// for decoding the file afterwards
	fwrite(&all_bits_written, sizeof(unsigned int), 1, fp_out);
}

void DecodeFile(FILE *fp_in, FILE *fp_out, DecodeHuffTree *decode_tree){
	unsigned int all_code_bits, bits_has_read;
	unsigned char current_byte;
	unsigned char check_bit_mask;
	char corresponding_symbol;
	void *ptTree;

	fseek(fp_in, 0 - (long)(sizeof(unsigned int)), SEEK_END);
	fread(&all_code_bits, sizeof(unsigned int), 1, fp_in);

	fseek(fp_in, 0, SEEK_SET);

	check_bit_mask = 0x80u;
	bits_has_read = 0;
	
	fread(&current_byte, sizeof(unsigned char), 1, fp_in);
	
	while(bits_has_read < all_code_bits){
		// writes one character in an iteration
		// of the outer loop

		ptTree = decode_tree;
		corresponding_symbol = -1;

		while(corresponding_symbol == -1){
			if(bits_has_read == all_code_bits){
				fprintf(stderr, "Length not match! The decoded file might not match the original.\n"
								"Make sure you get the right key and the right encoded file.\n");
				return;
			}

			if(current_byte & check_bit_mask)
				ptTree = ((DecodeHuffTreeInteriorNode *)ptTree)->rightChild;
			else
				ptTree = ((DecodeHuffTreeInteriorNode *)ptTree)->leftChild;
			
			bits_has_read++;

			if(DecodeHuffTreeIsLeaf(ptTree))
				corresponding_symbol = ((DecodeHuffTreeLeaf *)ptTree)->symbol;

			if(check_bit_mask == 0x01u){
				fread(&current_byte, sizeof(unsigned char), 1, fp_in);
				check_bit_mask = 0x80u;
			}else
				check_bit_mask >>= 1;
		}

		fputc(corresponding_symbol, fp_out);
	}

	// the file standards says that a non-empty text file
	// should end with a newline(0x0A).
	fputc(0x0A, fp_out);
}

void StartSamplingProcess(char *input_file_name, char *output_file_name){
	FILE *fp_in = fopen(input_file_name, "r");
	if(fp_in == NULL){
		fprintf(stderr, "Can not open file %s !\n", input_file_name);
		return;
	}
	FILE *fp_out = fopen(output_file_name, "wb");
	if(fp_out == NULL){
		fprintf(stderr, "Can not open file %s !\n", output_file_name);
		fclose(fp_in);
		return;
	}
	int i, freq_list[TOTAL_SYMBOL_NUMBER];

	GetFreqencyListFromSampleFile(fp_in, freq_list);

	for(i = 0; i < TOTAL_SYMBOL_NUMBER; i++)
		printf("%c(%d): %d\n", i + ASCII_PRINTABLE_CHAR_START, i + ASCII_PRINTABLE_CHAR_START, freq_list[i]);

	WriteFrequencyListToFile(fp_out, freq_list);
	
	fclose(fp_in);
	fclose(fp_out);
}

void StartEncodeProcess(char *input_file_name, char *output_file_name, char *key_file_name, int *freq_list){
	HuffmanForest F;
	HuffmanTree *T;
	EncodeTable encode_table;

	F = InitializeHuffmanForest(freq_list);
	T = ConstructHuffmanTree(F);
	encode_table = EncodeThroughHuffmanTree(T);

	PrintHuffTree(T);
	putchar('\n');
	ShowEncodeTable(encode_table);

	FILE *fp_in = fopen(input_file_name, "r");
	if(fp_in == NULL){
		fprintf(stderr, "Can not open file %s !\n", input_file_name);
		return;
	}
	FILE *fp_out = fopen(output_file_name, "wb");
	if(fp_out == NULL){
		fprintf(stderr, "Can not open file %s !\n", output_file_name);
		fclose(fp_in);
		return;
	}
	FILE *fp_k = fopen(key_file_name, "wb");
	if(fp_k == NULL){
		fprintf(stderr, "Can not open file %s !\n", key_file_name);
		fclose(fp_in);
		fclose(fp_out);
		return;
	}

	EncodeFile(fp_in, fp_out, encode_table);
	WriteDecodeHuffmanTreeToFile(T, fp_k);

	fclose(fp_in);
	fclose(fp_out);
	fclose(fp_k);

	DisposeHuffmanForest(F);
	DisposeHuffmanTree(T);
	DisposeEncodeTable(encode_table);
}

void StartDecodeProcess(char *input_file_name, char *output_file_name, char *key_file_name){
	// it turns out that there are differences
	// between 'r' and 'rb'
	FILE *fp_k = fopen(key_file_name, "rb");
	if(fp_k == NULL){
		fprintf(stderr, "Can not open file %s !\n", key_file_name);
		return;
	}
	FILE *fp_in = fopen(input_file_name, "rb");
	if(fp_in == NULL){
		fprintf(stderr, "Can not open file %s !\n", input_file_name);
		fclose(fp_k);
		return;
	}
	FILE *fp_out = fopen(output_file_name, "w");
	if(fp_out == NULL){
		fprintf(stderr, "Can not open file %s !\n", output_file_name);
		fclose(fp_in);
		fclose(fp_k);
		return;
	}
	DecodeHuffTree *decode_tree = ReadDecodeHuffmanTreeFromFile(fp_k);

	PrintDecodeHuffTree(decode_tree);
	putchar('\n');

	DecodeFile(fp_in, fp_out, decode_tree);

	fclose(fp_k);
	fclose(fp_in);
	fclose(fp_out);

	DisposeDecodeHuffTree(decode_tree);
}

#define __collect_parameter_input(loop_i, loop_j, argc, argv, collected_parameter, parameter_specific, parameter_associated_input, starts, end) \
	collected_parameter = 0;	\
	for(loop_i = 2; loop_i < argc; loop_i++){	\
		for(loop_j = starts; loop_j <= end; loop_j++){	\
			if(!strcmp(argv[loop_i], parameter_specific[loop_j])){	\
				if(loop_i + 1 < argc && argv[i + 1][0] != '-'){	\
					parameter_associated_input[loop_j] = argv[++i];	\
					collected_parameter++;	\
				}else{	\
					fprintf(stderr, "No file specified after %s!\n", parameter_specific[loop_j]);	\
					return;	\
				}	\
			}	\
		}	\
	}

void HandleParameters(int argc, char ** argv){
	int i, j, collected_parameter;
	char *parameter_specific[5] = {
		"--sample-file",
		"--sample-data",
		"-i",
		"-o",
		"-k"
	};
	char *parameter_description[5] = {
		"sample file",
		"sample data file",
		"input file",
		"output file",
		"key file"
	};
	char *parameter_associated_input[5] = {NULL};

	if(argc <= 2){
		fprintf(stdout, "A tool for huffman encoding/decoding.\n"
						"	Commands: \n"
						"\n"
						"	huffcodec sample [-i <file-name>] [-o <file-name>]:\n"
						"		Collect symbol appearance frequency data from input file, write the data\n"
						"		to output file, which can be used in the encoding process\n"
						"\n"
						"	huffcodec encode [--sample-from <file-name>]/[--sample-data <file-name>]\n"
						"	[-i <file-name>] [-o <file-name>] [-k <file-name>] :\n"
					    "		Encode symbols in input file with the information of symbol appearance\n"
					    "		frequency data either collected from the file specified by --sample-from\n" 
					    "		or read directly from the file specified by --sample-data. Write the en-\n"
					    "		coded information to output file, and the encoding information to the file\n"
					    "		specified by -k.\n"
						"\n"
					    "	huffcodec decode [-i <file-name>] [-o <file-name>] [-k <file-name>]:\n"
					    "		Decode the input file with the encoding information from the file specified\n"
					    "		by -k. Write the decoded information to output file.\n"
						"\n"
						"	huffcodec clean [--sample-file <file-name>]/[--sample-data <file-name>] [-i <file-name>]\n"
						"	[-o <file-name>]:\n"
						"		Clear all the symbols in input file that do not appear in sample file or have zero fre-\n"
						"		qencies in sample data, write the result into output file and print symbols that have\n"
						"		deleted.\n"
						"\n"
						"	NOTE: The relative orders of parameters of upper three commands do not matter.\n"
						"\n"
						"For more information, visit our repository at https://gitee.com/tinusgraglin/huffman-codec\n");
	}else{
		if(!strcmp(argv[1], "encode")){
			__collect_parameter_input(i, j, argc, argv, collected_parameter, parameter_specific, parameter_associated_input, 0, 4)

			if(parameter_associated_input[0] != NULL && parameter_associated_input[1] != NULL){
				fprintf(stderr, "Duplicated parameters: %s and %s!\n", parameter_specific[0], parameter_specific[1]);
				return;
			}else if(parameter_associated_input[0] == NULL && parameter_associated_input[1] == NULL){
				fprintf(stderr, "Please specify at least one of the parameters: %s and %s\n", parameter_specific[0], parameter_specific[1]);
				return;
			}

			for(j = 2; j < 5 && parameter_associated_input[j]; j++);
			if(j < 5){
				fprintf(stderr, "%s not specified!\n", parameter_description[j]);
				return;
			}

			for(j = 0; j < 5; j++)
				fprintf(stdout, "%s: %s\n", parameter_specific[j], (parameter_associated_input[j] == NULL)?"(Empty)":parameter_associated_input[j]);

			int freq_list[TOTAL_SYMBOL_NUMBER];
			FILE *fp;
			if(parameter_associated_input[0] == NULL){
				fp = fopen(parameter_associated_input[1], "rb");
				ReadFreqencyListFromFile(fp, freq_list);
			}else{
				fp = fopen(parameter_associated_input[0], "r");
				GetFreqencyListFromSampleFile(fp, freq_list);
			}
			fclose(fp);
			for(int i = 0; i < TOTAL_SYMBOL_NUMBER; i++)
				fprintf(stdout, "%c(%d): %d\n", i + ASCII_PRINTABLE_CHAR_START, i + ASCII_PRINTABLE_CHAR_START, freq_list[i]);
			StartEncodeProcess(parameter_associated_input[2], parameter_associated_input[3], parameter_associated_input[4], freq_list);

		}else if(!strcmp(argv[1], "decode")){
			__collect_parameter_input(i, j, argc, argv, collected_parameter, parameter_specific, parameter_associated_input, 2, 4)

			for(j = 2; j < 5 && parameter_associated_input[j]; j++);
			if(j < 5){
				fprintf(stderr, "%s not specified!\n", parameter_description[j]);
				return;
			}

			for(j = 2; j < 5; j++)
				fprintf(stdout, "%s: %s\n", parameter_specific[j], (parameter_associated_input[j] == NULL)?"(Empty)":parameter_associated_input[j]);

			StartDecodeProcess(parameter_associated_input[2], parameter_associated_input[3], parameter_associated_input[4]);

		}else if(!strcmp(argv[1], "sample")){
			__collect_parameter_input(i, j, argc, argv, collected_parameter, parameter_specific, parameter_associated_input, 2, 3)

			for(j = 2; j < 4 && parameter_associated_input[j]; j++);
			if(j < 4){
				fprintf(stderr, "%s not specified!\n", parameter_description[j]);
				return;
			}

			for(j = 2; j < 4; j++)
				fprintf(stdout, "%s: %s\n", parameter_specific[j], (parameter_associated_input[j] == NULL)?"(Empty)":parameter_associated_input[j]);

			StartSamplingProcess(parameter_associated_input[2], parameter_associated_input[3]);
		}else if(!strcmp(argv[1], "clean")){
			__collect_parameter_input(i, j, argc, argv, collected_parameter, parameter_specific, parameter_associated_input, 0, 3)

			if(parameter_associated_input[0] != NULL && parameter_associated_input[1] != NULL){
				fprintf(stderr, "Duplicated parameters: %s and %s!\n", parameter_specific[0], parameter_specific[1]);
				return;
			}else if(parameter_associated_input[0] == NULL && parameter_associated_input[1] == NULL){
				fprintf(stderr, "Please specify at least one of the parameters: %s and %s\n", parameter_specific[0], parameter_specific[1]);
				return;
			}

			for(j = 2; j < 4 && parameter_associated_input[j]; j++);
			if(j < 4){
				fprintf(stderr, "%s not specified!\n", parameter_description[j]);
				return;
			}

			for(j = 0; j < 4; j++)
				fprintf(stdout, "%s: %s\n", parameter_specific[j], (parameter_associated_input[j] == NULL)?"(Empty)":parameter_associated_input[j]);

			int freq_list[TOTAL_SYMBOL_NUMBER];
			FILE *fp_out, *fp, *fp_in;
			if(parameter_associated_input[0] != NULL){
				fp = fopen(parameter_associated_input[0], "r");
				GetFreqencyListFromSampleFile(fp, freq_list);
			}else{
				fp = fopen(parameter_associated_input[1], "rb");
				ReadFreqencyListFromFile(fp, freq_list);
			}
			fp_in = fopen(parameter_associated_input[2], "r");
			fp_out = fopen(parameter_associated_input[3], "w");
			ClearZeroFreqencySymbols(freq_list, fp_in, fp_out);
			fclose(fp);
			fclose(fp_in);
			fclose(fp_out);
		}
	}
}
